/*
 * topological-hex, a program to compute combinatorial hexahedral meshes.
 *
 * Copyright (C) <2018> <Université catholique de Louvain (UCL), Belgique>
 *
 * List of the contributors to the development, description and complete
 * License: see LICENSE file.
 *
 * This program (topological-hex) is free software:
 * you can redistribute it and/or modify it under the terms
 * of the GNU General Public License as published by the Free
 * Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program (see COPYING file).  If not,
 * see <http://www.gnu.org/licenses/>.
 */

#include "protocol.h"
#include "symmetry.h"

#include <string.h>
#include <stdlib.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <sys/select.h>
#include <netinet/tcp.h>
#include <unistd.h>
#include <fcntl.h>
#include <errno.h>

/**
 * Reads @p size bytes from @p socket into @p buffer.
 *
 * @return True upon success
 */
static
bool read_data(bool (*cond)(void *data), void *data,
               int socket, void *buffer, size_t size);

/**
 * Writes @p size bytes from @p buffer into @p socket.
 *
 * @return True upon success
 */
static
bool write_data(bool (*cond)(void *data), void *data,
                int socket, const void *buffer, size_t size);

error_code configure_nonblock(int socket) {
  int flags = fcntl(socket, F_GETFL, 0);
  if (flags == -1) return SYSTEM_ERROR;

  if (fcntl(socket, F_SETFL, flags | O_NONBLOCK) == -1)
    return SYSTEM_ERROR;

  return SUCCESS;
}

void configure_keepalive(int socket) {
  int one = 1;
  setsockopt(socket, SOL_SOCKET, SO_KEEPALIVE, &one, sizeof(one));

  int keep_count    = 5;
  int keep_idle     = 60;
  int keep_interval = 120;

  setsockopt(socket, SOL_SOCKET, TCP_KEEPCNT,
             &keep_count, sizeof(keep_count));
  setsockopt(socket, SOL_SOCKET, TCP_KEEPIDLE,
             &keep_idle, sizeof(keep_idle));
  setsockopt(socket, SOL_SOCKET, TCP_KEEPINTVL,
             &keep_interval, sizeof(keep_interval));
}

bool read_name(bool (*cond)(void *data), void *data,
               int socket, char *buffer, size_t buffer_len) {
  memset(buffer, 0, buffer_len);
  uint8_t len_buf[sizeof(uint64_t)];
  if (!read_data(cond, data, socket, len_buf, sizeof(len_buf)))
    return false;

  uint64_t len = decode_uint64(len_buf);
  if (len >= buffer_len) return false;
  return read_data(cond, data, socket, buffer, len);
}

bool write_name(bool (*cond)(void *data), void *data,
                int socket, const char *buffer, size_t len) {
  uint8_t len_buf[sizeof(uint64_t)];
  encode_uint64(len_buf, len);
  return write_data(cond, data, socket, len_buf, sizeof(len_buf)) &&
    write_data(cond, data, socket, buffer, len);
}

bool read_command(bool (*cond)(void *data), void *data,
                  int socket, message_tag *tag, uint64_t *arg) {
  uint8_t buffer[sizeof(uint64_t)];
  if (read_data(cond, data, socket, buffer, sizeof(buffer))) {
    uint64_t value = decode_uint64(buffer);

    *tag = value & ((1 << TAG_BITS)-1);
    *arg = value >> TAG_BITS;

    return true;
  }
  else
    return false;
}

bool write_command(bool (*cond)(void *data), void *data,
                   int socket, message_tag tag, uint64_t arg) {
  uint64_t value = ((uint64_t)arg << (uint64_t)TAG_BITS) | tag;
  uint8_t buffer[sizeof(uint64_t)];
  encode_uint64(buffer, value);
  return write_data(cond, data, socket, buffer, sizeof(buffer));
}

bool read_layout(bool (*cond)(void *data), void *data,
                 int socket,
                 uint32_t *num_vertices, uint32_t *num_quads,
                 uint32_t *num_hexes, uint32_t *num_symmetries,
                 uint32_t *num_boundary_vertices,
                 uint32_t *max_no_goods_size, uint8_t *flags) {
  uint8_t buffer[7*sizeof(uint32_t)];
  if (read_data(cond, data, socket, buffer, sizeof(buffer))) {
    *num_vertices = decode_uint32(buffer);
    *num_quads = decode_uint32(buffer + sizeof(uint32_t));
    *num_hexes = decode_uint32(buffer + 2*sizeof(uint32_t));
    *num_symmetries = decode_uint32(buffer + 3*sizeof(uint32_t));
    *num_boundary_vertices = decode_uint32(buffer + 4*sizeof(uint32_t));
    *max_no_goods_size = decode_uint32(buffer + 5*sizeof(uint32_t));
    *flags = buffer[6*sizeof(uint32_t)];
    return true;
  }
  else
    return false;
}

bool write_layout(bool (*cond)(void *data), void *data,
                  int socket,
                  uint32_t num_vertices, uint32_t num_quads,
                  uint32_t num_hexes, uint32_t num_symmetries,
                  uint32_t num_boundary_vertices,
                  uint32_t max_no_goods_size, uint8_t flags) {
  uint8_t buffer[7*sizeof(uint32_t)];
  encode_uint32(buffer, num_vertices);
  encode_uint32(buffer + sizeof(uint32_t), num_quads);
  encode_uint32(buffer + 2*sizeof(uint32_t), num_hexes);
  encode_uint32(buffer + 3*sizeof(uint32_t), num_symmetries);
  encode_uint32(buffer + 4*sizeof(uint32_t), num_boundary_vertices);
  encode_uint32(buffer + 5*sizeof(uint32_t), max_no_goods_size);
  buffer[6*sizeof(uint32_t)] = flags;
  return write_data(cond, data, socket, buffer, sizeof(buffer));
}

bool read_solution(bool (*cond)(void *data), void *data,
                   int socket,
                   uint8_t *buffer,
                   vertex_index **vertices,
                   hex_index *num_hexes, uint32_t max_hexes) {
  if (!read_data(cond, data, socket, buffer, sizeof(uint8_t)))
    return false;

  *num_hexes = buffer[0];

  if (*num_hexes > max_hexes)
    return false;

  if (!read_data(cond, data, socket, buffer, *num_hexes * sizeof(uint8_t) * 8))
    return false;

  *vertices = (vertex_index*)buffer;

  return true;
}

bool write_solution(bool (*cond)(void *data), void *data,
                    int socket, uint8_t *buffer,
                    const vertex_index *vertices, hex_index num_hexes) {
  if (!write_data(cond, data, socket, &num_hexes, sizeof(num_hexes)))
    return false;

  return write_data(cond, data, socket, vertices,
                    (num_hexes * 8) * sizeof(uint8_t));
}

size_t compute_subproblem_layout(
  subproblem_layout *layout,
  uint32_t num_vertices, uint32_t num_quads, uint32_t num_hexes,
  uint32_t num_symmetries, uint32_t num_boundary_vertices,
  uint32_t max_no_goods_size, uint8_t flags) {
  layout->num_vertices          = num_vertices;
  layout->num_quads             = num_quads;
  layout->num_hexes             = num_hexes;
  layout->num_symmetries        = num_symmetries;
  layout->num_boundary_vertices = num_boundary_vertices;
  layout->max_no_goods_size     = max_no_goods_size;
  layout->flags                 = flags;

  size_t offset = 0;
  layout->offset_hex_id = offset; offset += sizeof(uint8_t);

  layout->offset_max_quad_id       = offset; offset += sizeof(uint32_t);
  layout->offset_max_vertex_id     = offset; offset += sizeof(uint32_t);
  layout->offset_quad_queue_size   = offset; offset += sizeof(uint32_t);
  layout->offset_quad_queue_offset = offset; offset += sizeof(uint32_t);
  layout->offset_bound_mask        = offset; offset += sizeof(uint32_t);

  layout->offset_quads = offset; offset += 4 * sizeof(uint8_t) * num_quads;
  layout->offset_hexes = offset; offset += 8 * sizeof(uint8_t) * num_hexes;

  layout->offset_quad_adjacency = offset;
  if (flags & SOLVER_SHELLABLE_ONLY)
    offset += 4 * sizeof(uint16_t) * num_quads;

  layout->offset_quad_queue = offset;
  offset += sizeof(uint8_t) * num_quads;

  layout->offset_num_occurrences = offset;
  offset += sizeof(uint8_t) * num_quads;

  layout->offset_quad_descriptors = offset;
  offset += sizeof(uint32_t) * num_vertices * num_vertices;

  layout->offset_allowed_neighbors = offset;
  offset += VERTEX_SET_NUM_WORDS * sizeof(uint64_t) * num_vertices;

  layout->offset_known_neighbors = offset;
  offset += VERTEX_SET_NUM_WORDS * sizeof(uint64_t) * num_vertices;

  layout->offset_known_diagonals = offset;
  offset += VERTEX_SET_NUM_WORDS * sizeof(uint64_t) * num_vertices;

  layout->offset_known_quad_diagonals = offset;
  offset += VERTEX_SET_NUM_WORDS * sizeof(uint64_t) * num_vertices;

  layout->offset_domains = offset;
  offset += VERTEX_SET_NUM_WORDS * sizeof(uint64_t) * 4;

  layout->offset_even_vertices = offset;
  offset += VERTEX_SET_NUM_WORDS * sizeof(uint64_t);

  layout->offset_odd_vertices = offset;
  offset += VERTEX_SET_NUM_WORDS * sizeof(uint64_t);

  layout->offset_boundary_group = offset;
  offset += sizeof(vertex_index) * num_boundary_vertices * num_symmetries;

  layout->offset_nogoods_size = offset;
  offset += sizeof(uint64_t);

  layout->offset_nogoods = offset;
  offset += sizeof(vertex_index) * max_no_goods_size;

  layout->offset_symmetry_data = offset;
  offset += sizeof(uint8_t) * solver_incremental_symmetry_data_size(
    num_vertices, num_hexes, num_symmetries);

  layout->packet_size = offset;
  return offset;
}

bool read_raw_subproblem(bool (*cond)(void *data), void *data,
                         int socket, uint8_t *buffer,
                         const subproblem_layout *layout) {
  return read_data(cond, data, socket, buffer, layout->packet_size);
}

bool write_raw_subproblem(bool (*cond)(void *data), void *data,
                          int socket, const uint8_t *buffer,
                          const subproblem_layout *layout) {
  return write_data(cond, data, socket, buffer, layout->packet_size);
}

bool read_subproblem(bool (*cond)(void *data), void *data,
                     int socket,
                     uint8_t *buffer, const subproblem_layout *layout,
                     solver *solver) {
  if (!read_data(cond, data, socket, buffer, layout->packet_size))
    return false;

  solver->hex_id = buffer[layout->offset_hex_id];

  solver->max_quad_id   = decode_uint32(&buffer[layout->offset_max_quad_id]);
  solver->max_vertex_id = decode_uint32(&buffer[layout->offset_max_vertex_id]);

  solver->quad_queue_offset =
    decode_uint32(&buffer[layout->offset_quad_queue_offset]);
  solver->quad_queue_size =
    decode_uint32(&buffer[layout->offset_quad_queue_size]);

  solver->bound_mask = decode_uint32(&buffer[layout->offset_bound_mask]);

  for (uint32_t i = 0; i < 4 * layout->num_quads; i++)
    solver->quads[i] = buffer[layout->offset_quads + i];

  for (uint32_t i = 0; i < 8 * layout->num_hexes; i++)
    solver->hexes[i] = buffer[layout->offset_hexes + i];

  if (layout->flags & SOLVER_SHELLABLE_ONLY) {
    for (uint32_t i = 0; i < 4 * layout->num_quads; i++) {
      solver->quad_adjacency[i] =
        decode_uint16(buffer + layout->offset_quad_adjacency +
                      i * sizeof(uint16_t));
    }
  }

  for (uint32_t i = 0; i < layout->num_quads; i++)
    solver->quad_queue[i] = buffer[layout->offset_quad_queue + i];

  for (uint32_t i = 0; i < layout->num_quads; i++)
    solver->num_occurrences[i] = buffer[layout->offset_num_occurrences + i];

  for (uint32_t i = 0; i < layout->num_vertices * layout->num_vertices; i++) {
    solver->quad_descriptors[i] =
      decode_uint32(buffer + layout->offset_quad_descriptors +
                    i*sizeof(uint32_t));
  }

  for (uint32_t i = 0; i < solver->num_vertices; i++) {
    for (uint32_t j = 0; j < VERTEX_SET_NUM_WORDS; j++) {
      solver->allowed_neighbors[i].data[j] =
        decode_uint64(buffer + layout->offset_allowed_neighbors +
                      (VERTEX_SET_NUM_WORDS * i + j) * sizeof(uint64_t));
    }
  }

  for (uint32_t i = 0; i < solver->num_vertices; i++) {
    for (uint32_t j = 0; j < VERTEX_SET_NUM_WORDS; j++) {
      solver->known_neighbors[i].data[j] =
        decode_uint64(buffer + layout->offset_known_neighbors +
                      (VERTEX_SET_NUM_WORDS * i + j) * sizeof(uint64_t));
    }
  }

  for (uint32_t i = 0; i < solver->num_vertices; i++) {
    for (uint32_t j = 0; j < VERTEX_SET_NUM_WORDS; j++) {
      solver->known_diagonals[i].data[j] =
        decode_uint64(buffer + layout->offset_known_diagonals +
                      (VERTEX_SET_NUM_WORDS * i + j) * sizeof(uint64_t));
    }
  }

  for (uint32_t i = 0; i < solver->num_vertices; i++) {
    for (uint32_t j = 0; j < VERTEX_SET_NUM_WORDS; j++) {
      solver->known_quad_diagonals[i].data[j] =
        decode_uint64(buffer + layout->offset_known_quad_diagonals +
                      (VERTEX_SET_NUM_WORDS * i + j) * sizeof(uint64_t));
    }
  }

  for (uint32_t i = 0; i < 4; i++) {
    for (uint32_t j = 0; j < VERTEX_SET_NUM_WORDS; j++) {
      solver->domains[i].data[j] =
        decode_uint64(buffer + layout->offset_domains +
                      (VERTEX_SET_NUM_WORDS * i + j) * sizeof(uint64_t));
    }
  }

  for (uint32_t i = 0; i < VERTEX_SET_NUM_WORDS; i++) {
    solver->even_vertices.data[i] =
      decode_uint64(buffer + layout->offset_even_vertices + i * sizeof(uint64_t));
  }

  for (uint32_t i = 0; i < VERTEX_SET_NUM_WORDS; i++) {
    solver->odd_vertices.data[i] =
      decode_uint64(buffer + layout->offset_odd_vertices + i * sizeof(uint64_t));
  }

  if (solver->num_symmetries != 0) {
    memcpy(solver->boundary_group, buffer + layout->offset_boundary_group,
           sizeof(*solver->boundary_group) *
           solver->num_boundary_vertices * solver->num_symmetries);

    if (solver->flags & SOLVER_SHELLABLE_ONLY)
      memcpy(solver->symmetry_data, buffer + layout->offset_symmetry_data,
             solver_incremental_symmetry_data_size(
               layout->num_vertices, layout->num_hexes, layout->num_symmetries));
  }

  solver->no_goods_size = decode_uint64(buffer + layout->offset_nogoods_size);

  if (solver->no_goods_size != 0) {
    if (solver->no_goods_capacity < solver->no_goods_size) {
      free(solver->no_goods);
      solver->no_goods = NULL;
      solver->no_goods_capacity = 0;

      if (posix_memalign((void**)&solver->no_goods, ALIGNMENT,
                         sizeof(*solver->no_goods) *
                         solver->no_goods_size) != 0)
        return false;

      solver->no_goods_capacity = solver->no_goods_size;
    }

    memcpy(solver->no_goods, buffer + layout->offset_nogoods,
           solver->no_goods_size * sizeof(*solver->no_goods));
  }

  return true;
}

bool write_subproblem(bool (*cond)(void *data), void *data,
                      int socket,
                      uint8_t *buffer,
                      const subproblem_layout *layout,
                      solver *solver) {
  buffer[layout->offset_hex_id] = solver->hex_id;

  encode_uint32(&buffer[layout->offset_max_quad_id], solver->max_quad_id);
  encode_uint32(&buffer[layout->offset_max_vertex_id], solver->max_vertex_id);

  encode_uint32(&buffer[layout->offset_quad_queue_offset],
                solver->quad_queue_offset);
  encode_uint32(&buffer[layout->offset_quad_queue_size],
                solver->quad_queue_size);

  encode_uint32(&buffer[layout->offset_bound_mask], solver->bound_mask);

  for (uint32_t i = 0; i < 4 * layout->num_quads; i++)
    buffer[layout->offset_quads + i] = solver->quads[i];

  for (uint32_t i = 0; i < 8 * layout->num_hexes; i++)
    buffer[layout->offset_hexes + i] = solver->hexes[i];

  if (layout->flags & SOLVER_SHELLABLE_ONLY) {
    for (uint32_t i = 0; i < 4 * layout->num_quads; i++) {
      encode_uint16(
        buffer + layout->offset_quad_adjacency + i * sizeof(uint16_t),
        solver->quad_adjacency[i]);
    }
  }

  for (uint32_t i = 0; i < layout->num_quads; i++)
    buffer[layout->offset_quad_queue + i] = solver->quad_queue[i];

  for (uint32_t i = 0; i < layout->num_quads; i++)
    buffer[layout->offset_num_occurrences + i] = solver->num_occurrences[i];

  for (uint32_t i = 0; i < layout->num_vertices * layout->num_vertices; i++) {
    encode_uint32(
      buffer + layout->offset_quad_descriptors +
      i*sizeof(uint32_t), solver->quad_descriptors[i]);
  }

  for (uint32_t i = 0; i < solver->num_vertices; i++) {
    for (uint32_t j = 0; j < VERTEX_SET_NUM_WORDS; j++) {
      encode_uint64(
        buffer + layout->offset_allowed_neighbors +
        (VERTEX_SET_NUM_WORDS * i + j) * sizeof(uint64_t),
        solver->allowed_neighbors[i].data[j]);
    }
  }

  for (uint32_t i = 0; i < solver->num_vertices; i++) {
    for (uint32_t j = 0; j < VERTEX_SET_NUM_WORDS; j++) {
      encode_uint64(
        buffer + layout->offset_known_neighbors +
        (VERTEX_SET_NUM_WORDS * i + j) * sizeof(uint64_t),
        solver->known_neighbors[i].data[j]);
    }
  }

  for (uint32_t i = 0; i < solver->num_vertices; i++) {
    for (uint32_t j = 0; j < VERTEX_SET_NUM_WORDS; j++) {
      encode_uint64(
        buffer + layout->offset_known_diagonals +
        (VERTEX_SET_NUM_WORDS * i + j) * sizeof(uint64_t),
        solver->known_diagonals[i].data[j]);
    }
  }

  for (uint32_t i = 0; i < solver->num_vertices; i++) {
    for (uint32_t j = 0; j < VERTEX_SET_NUM_WORDS; j++) {
      encode_uint64(
        buffer + layout->offset_known_diagonals +
        (VERTEX_SET_NUM_WORDS * i + j) * sizeof(uint64_t),
        solver->known_diagonals[i].data[j]);
    }
  }

  for (uint32_t i = 0; i < 4; i++) {
    for (uint32_t j = 0; j < VERTEX_SET_NUM_WORDS; j++) {
      encode_uint64(
        buffer + layout->offset_domains +
        (VERTEX_SET_NUM_WORDS * i + j) * sizeof(uint64_t),
        solver->domains[i].data[j]);
    }
  }

  for (uint32_t i = 0; i < VERTEX_SET_NUM_WORDS; i++) {
    encode_uint64(buffer + layout->offset_even_vertices + i * sizeof(uint64_t),
                  solver->even_vertices.data[i]);
  }

  for (uint32_t i = 0; i < VERTEX_SET_NUM_WORDS; i++) {
    encode_uint64(buffer + layout->offset_odd_vertices + i * sizeof(uint64_t),
                  solver->odd_vertices.data[i]);
  }

  if (solver->num_symmetries != 0) {
    memcpy(buffer + layout->offset_boundary_group, solver->boundary_group,
           sizeof(*solver->boundary_group) *
           solver->num_boundary_vertices * solver->num_symmetries);

    if (solver->flags & SOLVER_SHELLABLE_ONLY)
      memcpy(buffer + layout->offset_symmetry_data, solver->symmetry_data,
             solver_incremental_symmetry_data_size(
               layout->num_vertices, layout->num_hexes, layout->num_symmetries));
  }

  encode_uint64(buffer + layout->offset_nogoods_size,
                solver->no_goods_size);

  if (solver->no_goods_size != 0) {
    memcpy(buffer + layout->offset_nogoods, solver->no_goods,
           solver->no_goods_size * sizeof(*solver->no_goods));
  }

  return write_data(cond, data, socket, buffer, layout->packet_size);
}

void encode_uint64(void *buffer, uint64_t x) {
  for (size_t i = 0; i < sizeof(x); i++)
    ((uint8_t*)buffer)[i] = (x >> (8*i)) & UINT8_MAX;
}

uint64_t decode_uint64(void *buffer) {
  uint64_t x = 0;
  for (size_t i = 0; i < sizeof(x); i++)
    x |= ((uint64_t)((uint8_t*)buffer)[i]) << (uint64_t)(8*i);
  return x;
}

void encode_uint32(void *buffer, uint32_t x) {
  for (size_t i = 0; i < sizeof(x); i++)
    ((uint8_t*)buffer)[i] = (x >> (8*i)) & UINT8_MAX;
}

uint32_t decode_uint32(void *buffer) {
  uint32_t x = 0;
  for (size_t i = 0; i < sizeof(x); i++)
    x |= ((uint32_t)((uint8_t*)buffer)[i]) << (uint32_t)(8*i);
  return x;
}

void encode_uint16(void *buffer, uint16_t x) {
  for (size_t i = 0; i < sizeof(x); i++)
    ((uint8_t*)buffer)[i] = (x >> (8*i)) & UINT8_MAX;
}

uint16_t decode_uint16(void *buffer) {
  uint16_t x = 0;
  for (size_t i = 0; i < sizeof(x); i++)
    x |= ((uint16_t)((uint8_t*)buffer)[i]) << (uint16_t)(8*i);
  return x;
}

static
bool read_data(bool (*cond)(void *data), void *data,
               int socket, void *buffer, size_t size) {
  size_t offset = 0;

  while (offset < size && cond(data)) {
    fd_set to_read;
    FD_ZERO(&to_read);
    FD_SET(socket, &to_read);

    struct timeval timeout;
    timeout.tv_sec  = 1;
    timeout.tv_usec = 0;

    int result = select(socket + 1, &to_read, NULL, NULL, &timeout);
    if (result == -1)
      return false;
    else if (result == 0)
      continue;

    ssize_t num_read = read(socket, (char*)buffer + offset, size - offset);

    if (num_read == -1) {
      if (errno == EAGAIN) continue;
      else return false;
    }
    else if (num_read == 0)
      return false;

    offset += num_read;
  }

  return offset == size;
}

static
bool write_data(bool (*cond)(void *data), void *data,
                int socket, const void *buffer, size_t size) {
  size_t offset = 0;

  while (offset < size && cond(data)) {
    fd_set to_write;

    FD_ZERO(&to_write);
    FD_SET(socket, &to_write);

    struct timeval timeout;
    timeout.tv_sec  = 1;
    timeout.tv_usec = 0;

    int result = select(socket + 1, NULL, &to_write, NULL, &timeout);
    if (result == -1)
      return false;
    else if (result == 0)
      continue;

    if (FD_ISSET(socket, &to_write)) {
      ssize_t num_write = write(socket, (char*)buffer + offset, size - offset);

      if (num_write == -1) {
        if (errno == EAGAIN) continue;
        else return false;
      }

      offset += num_write;
    }
  }

  return offset == size;
}
